#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""实现自定义WHOIS查询功能的模块。

该模块封装了WHOIS查询的相关操作，支持从不同的WHOIS服务器获取域名信息。 

Attributes:
    __version__ (str): "0.0.1"
    __copyright__ (str): "Copyright (C) 2025 z"
    __license__ (str): "MIT"
    __author__ (str): "Z"
"""

import io
import re
import socket
from copy import deepcopy
from typing import Dict, Optional, Union

import requests

from wingman import logger
from wingman.common.error import ErrorCode
from wingman import epochExpire

# 根服务器地址
ROOT_SERVER = 'whois.iana.org'

# 默认的配置，适用于大部分域名查询
DEFAULT_WHOIS_CONFIG = {
    'whois_server': '',  # whois查询服务器
    'registry_time': 'Creation Date',  # 注册时间
    'expire_time': 'Registry Expiry Date',  # 过期时间
    'registrar': 'Registrar',  # 注册商
    'registrar_url': 'Registrar URL',  # 注册商url
}

# 自定义配置
CUSTOM_WHOIS_CONFIGS = {
    'cn': {
        'whois_server': 'whois.cnnic.cn',
        'registry_time': 'Registration Time',
        'expire_time': 'Expiration Time',
        'registrar': 'Sponsoring Registrar',
    },
    'hk': {
        'whois_server': 'whois.hkirc.hk',
        'registry_time': 'Domain Name Commencement Date',
        "registry_time_format": '%d-%m-%Y',
        'expire_time': 'Expiry Date',
        "expire_time_format": '%d-%m-%Y',
    },
    'sg': {
        'whois_server': 'whois.sgnic.sg',
        'registry_time': 'Creation Date',
        'expire_time': 'Expiration Date',
    },
    'jp': {
        'whois_server': 'whois.jprs.jp',
        'registry_time': '[登録年月日]',
        "registry_time_format": '%Y/%m/%d',
        'expire_time': '[有効期限]',
        "expire_time_format": '%Y/%m/%d',
    },

    '中国': {
        'whois_server': 'cwhois.cnnic.cn',
        'registry_time': 'Registration Time',
        'expire_time': 'Expiration Time',
    },
    'tw': {
        # 'whois_server': 'whois.twnic.net',
        'registry_time': 'Record created',
        "registry_time_format": '%Y-%m-%d %H:%M:%S (UTC+8)',
        'expire_time': 'Record expires',
        "expire_time_format": '%Y-%m-%d %H:%M:%S (UTC+8)',
    },
    'ws': {
        'registry_time': 'Creation Date',
        'expire_time': 'Registrar Registration Expiration Date',
    },
    # 网页查询：https://www.whois.cm/
    'cm': {
        'whois_server': 'whois.registrar.cm',
        'registry_time': 'Creation Date',
        'expire_time': 'Registry Expiry Date'
    },
    'by': {
        'whois_server': 'whois.cctld.by',
        'registry_time': 'Creation Date',
        'expire_time': 'Expiration Date',
    },
    'im': {
        'expire_time': 'Expiry Date',
        'expire_time_format': '%d/%m/%Y %H:%M:%S',
    },
    'kr': {
        'whois_server': 'whois.kr',
        'registry_time': 'Registered Date',
        "registry_time_format": '%Y. %m. %d.',
        'expire_time': 'Expiration Date',
        "expire_time_format": '%Y. %m. %d.',
    },
}

# 国内cn域名注册商
REGISTRAR_CONFIG = [
    {
        'registrar': '厦门易名科技股份有限公司',
        'registrar_url': 'https://www.ename.net/'
    },
    {
        'registrar': '北京中科三方网络技术有限公司',
        'registrar_url': 'https://www.sfn.cn/'
    }
]

# 国内cn域名注册商map
REGISTRAR_CONFIG_MAP = {
    config['registrar']: config
    for config in REGISTRAR_CONFIG
}


class CustomWhois:
    """实现自定义WHOIS查询功能的类。

    该类封装了WHOIS查询的相关操作，支持从不同的WHOIS服务器获取域名信息。

    References:
        - 国际WHOIS组织: https://www.whois.com/whois/
        - IANA根域名数据库: https://www.iana.org/domains/root/db

    Note:
        该类主要用于处理国际域名的WHOIS查询，支持多种顶级域名（TLD）。
    """

    def __encode_domain(self, domain: str) -> str:
        """将域名转换为IDNA编码格式。

        该方法用于处理包含非ASCII字符（如中文）的域名，将其转换为IDNA（国际化域名）编码格式。
        对于纯ASCII域名（如英文域名），将原样返回。

        IDNA（Internationalizing Domain Names in Applications）是一种用于处理国际化域名的标准，
        它允许在域名中使用非ASCII字符，同时保持与现有DNS系统的兼容性。

        Args:
            domain (str): 要编码的域名，可以是中文或英文域名。
                示例："中文.com" 或 "example.com"

        Returns:
            str: 编码后的域名。对于中文域名，返回IDNA编码格式（如"xn--fiq228c.com"）；
                对于纯ASCII域名，返回原字符串。

        Examples:
            >>> __encode_domain("中文.com")
            'xn--fiq228c.com'
            >>> __encode_domain("example.com")
            'example.com'
        """
        return domain.encode('idna').decode('ascii')

    def __get_whois_raw(
        self,
        domain: str,
        server: str,
        timeout: int = 5,
        port: int = 43
    ) -> str:
        """请求WHOIS服务获取域名的原始信息。

        该方法会向指定的WHOIS服务器发送请求，获取域名的完整WHOIS信息。

        Args:
            domain: 要查询的域名，例如"example.com"。
            server: WHOIS服务器地址，例如"whois.verisign-grs.com"。
            timeout: 请求超时时间，默认5秒。
            port: WHOIS服务器端口，默认43。

        Returns:
            包含WHOIS原始信息的字符串，格式如下：
            Domain Name: LNJINQI.COM
            Registry Domain ID: 2865460509_DOMAIN_COM-VRSN
            Registrar WHOIS Server: grs-whois.hichina.com
            Registrar URL: http://wanwang.aliyun.com
            Updated Date: 2025-03-11T01:29:09Z
            Creation Date: 2024-03-21T07:38:45Z
            Registry Expiry Date: 2027-03-21T07:38:45Z
            ......

        Raises:
            ConnectionError: 如果无法连接到WHOIS服务器。
            TimeoutError: 如果请求超时。

        Example:
            >>> raw_info = __get_whois_raw("example.com", "whois.verisign-grs.com")
            >>> print(raw_info[:100])  # 打印前100个字符
        """
        # 创建连接
        sock = socket.create_connection((server, port))
        sock.settimeout(timeout)
        # 发送请求
        sock.send(("%s\r\n" % domain).encode("utf-8"))
        # 接收数据
        buff = bytes()
        while True:
            data = sock.recv(1024)
            if len(data) == 0:
                break
            buff += data
        # 关闭链接
        sock.close()
        return buff.decode("utf-8")

    def __load_whois_servers_txt(self) -> Dict[str, str]:
        """加载并返回whois服务器配置。

        该方法会从文本文件中加载whois服务器配置，返回一个将顶级域名（TLD）映射到对应whois服务器地址的字典。

        Returns:
            Dict[str, str]: 包含whois服务器配置的字典，格式为{
                'com': 'whois.verisign-grs.com',
                'cn': 'whois.cnnic.cn',
                ...
            }

        Raises:
            FileNotFoundError: 如果配置文件不存在。
            ValueError: 如果配置文件格式不正确。

        Example:
            >>> config = __load_whois_servers_txt()
            >>> print(config['com'])
            'whois.verisign-grs.com'
        """
        content_dict = {}
        file_path = "/app/wingman/domain/whois-servers.txt"
        with io.open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                split_line = line.strip().split(maxsplit=1)
                if len(split_line) == 2:
                    root, server = split_line
                    content_dict[root.strip()] = server.strip()
        # logger.debug(ErrorCode.SUCCESS.value, content_dist, __name__)
        return content_dict

    def __load_whois_servers_config(self) -> Dict[str, str]:
        """加载并返回whois服务器配置。

        该方法会加载whois服务器配置，返回一个将顶级域名（TLD）映射到对应whois服务器地址的字典。

        Returns:
            Dict[str, str]: 包含whois服务器配置的字典，格式为{
                'com': 'whois.verisign-grs.com',
                'cn': 'whois.cnnic.cn',
                ...
            }

        Raises:
            FileNotFoundError: 如果配置文件不存在。
            ValueError: 如果配置文件格式不正确。

        Example:
            >>> config = __load_whois_servers_config()
            >>> print(config['com'])
            'whois.verisign-grs.com'
        """
        whois_servers = self.__load_whois_servers_txt()
        config_dict = {}
        # 通用配置
        for root, server in whois_servers.items():
            server_config = deepcopy(DEFAULT_WHOIS_CONFIG)
            server_config['whois_server'] = server
            config_dict[self.__encode_domain(root)] = server_config
        # 自定义配置优先
        for key, value in CUSTOM_WHOIS_CONFIGS.items():
            encode_key = self.__encode_domain(key)
            default_config = config_dict.get(encode_key, deepcopy(DEFAULT_WHOIS_CONFIG))
            default_config.update(value)
            config_dict[encode_key] = default_config
            # 合并配置
        return config_dict

    # TODO: 实现从根服务器获取域名的查询服务器 (负责人: Z, 截止日期: 2023-12-31)
    # 该方法需要从根服务器（ROOT_SERVER）获取域名的查询服务器地址。
    # 目前未实现，计划使用正则表达式从原始数据中提取 "refer:" 字段。
    def __get_domain_root_whois(self, domain: str) -> str:
        raw_data = self.__get_whois_raw(domain, ROOT_SERVER, timeout=10)
        result = re.findall("refer:(.*)", raw_data)
        if result and len(result) > 0:
            return result[0].strip()

    def __get_whois_config(self, domain: str) -> str:
        """获取指定域名的一级域信息所在服务器。

        该方法会根据域名的一级域（TLD）查找对应的WHOIS服务器地址。

        Args:
            domain: 要查询的域名，例如"x.com"。

        Returns:
            对应一级域的WHOIS服务器地址，例如"whois.verisign-grs.com"。

        Raises:
            ValueError: 如果无法找到对应的一级域配置。

        Example:
            >>> __get_whois_config("example.com")
            'whois.verisign-grs.com'
        """

        domain_root = domain.split('.')[-1]
        WHOIS_CONFIGS = self.__load_whois_servers_config()

        if domain_root in WHOIS_CONFIGS:
            return WHOIS_CONFIGS.get(domain_root)
        # TODO: 实现从根服务器获取域名的查询服务器对应部分 (负责人: Z, 截止日期: 2025-12-31)
        # 从根服务器查询域名信息服务器
        else:
            domain_whois_server = self.__get_domain_root_whois(domain)
            if domain_whois_server:
                server_config = deepcopy(DEFAULT_WHOIS_CONFIG)
                server_config['whois_server'] = domain_whois_server
                return server_config
            else:
                raise Exception('not support {}'.format(domain))

    def get_domain_raw_whois(self, domain: str) -> str:
        """获取域名的WHOIS原始信息。

        该方法会查询指定域名的WHOIS信息，并返回完整的原始数据。

        Args:
            domain: 要查询的域名，例如"x.com"。

        Returns:
            包含WHOIS原始信息的字符串，格式如下：
            Domain Name: X.COM
            Registry Domain ID: 2865460509_DOMAIN_COM-VRSN
            Registrar WHOIS Server: grs-whois.hichina.com
            Registrar URL: http://wanwang.aliyun.com
            Updated Date: 2025-03-11T01:29:09Z
            Creation Date: 2024-03-21T07:38:45Z
            Registry Expiry Date: 2027-03-21T07:38:45Z
            ......

        Raises:
            ConnectionError: 如果无法连接到WHOIS服务器。
            TimeoutError: 如果查询超时。

        Example:
            >>> raw_whois = get_domain_raw_whois("example.com")
            >>> print(raw_whois[:100])  # 打印前100个字符
        """

        whois_config = self.__get_whois_config(domain)
        whois_server = whois_config['whois_server']
        raw_data = self.__get_whois_raw(domain, whois_server, timeout=10)
        # logger.debug(ErrorCode.SUCCESS.value, raw_data, __name__)
        return raw_data

    def __parse_whois_raw(self, whois_raw: str) -> Dict[str, str]:
        """解析WHOIS原始信息为键值对字典。

        该方法用于将WHOIS查询的原始信息解析为结构化的字典格式，
        为get_domain_whois函数提供数据支持。

        Args:
            whois_raw: 从WHOIS查询获取的原始信息字符串。

        Returns:
            包含解析结果的字典，格式为{
                'whois_server': str,
                'registry_time': str,
                'expire_time': str,
                'registrar': str,
                'registrar_url': str
            }

        Example:
            >>> raw_data = "Creation Date: 2021-01-01\nRegistrar: Example Inc."
            >>> __parse_whois_raw(raw_data)
            {
                'whois_server': 'whois.verisign-grs.com',
                'registry_time': '2021-01-01',
                'expire_time': '',
                'registrar': 'Example Inc.',
                'registrar_url': ''
            }
        """

        content = {}
        for row in whois_raw.split("\n"):
            # tw
            if 'Record expires on' in row or 'Record created on' in row:
                row_split = row.split("on", maxsplit=1)
            elif ":" in row:
                # fix: Python2 split() takes no keyword arguments
                row_split = row.split(":", 1)
            else:
                row_split = row.split(" ", 1)

            if len(row_split) == 2:
                key, value = row_split
                content[key.strip()] = value.strip()
        return content

    def get_domain_whois(self, domain: str) -> Optional[Dict[str, Union[str, None]]]:
        """提供字典形式的WHOIS内容

        Args:
            domain: 要查询的域名，如"x.com"

        Returns:
            包含WHOIS信息的字典，格式为{
                'domain': str,
                'start_time': str,
                'registrar': str,
                'registrar_url': str,
                'expire_time': str,
                'whois_server': str
            } 或 None（当无法获取有效信息时）

        Examples:
            >>> get_domain_whois("example.com")
            {
                'domain': 'example.com',
                'start_time': '2021-01-08T09:59:51Z',
                'registrar': 'Example Registrar',
                'registrar_url': 'http://example.com',
                'expire_time': '2031-01-08T09:59:51Z',
                'whois_server': 'whois.example.com'
            }
        """
        content = self.__parse_whois_raw(self.get_domain_raw_whois(domain))
        whois_config = self.__get_whois_config(domain)

        start_time = content.get(whois_config['registry_time'], '')
        expire_time = content.get(whois_config['expire_time'], '')
        registrar = content.get(whois_config.get('registrar'), '').strip()
        registrar_url_key = whois_config.get('registrar_url')
        registrar_url = content.get(registrar_url_key, '').strip()

        # cn域名注册商
        if registrar and not registrar_url:
            registrar_config = REGISTRAR_CONFIG_MAP.get(registrar)
            if registrar_config:
                registrar_url = registrar_config['registrar_url']

        if start_time or expire_time:
            return {
                'domain': domain,
                'start_time': start_time,
                'registrar': registrar,
                'registrar_url': registrar_url,
                'expire_time': expire_time,
                'whois_server': whois_config['whois_server'],
            }
        else:
            return None

    def get_domain_expire(self, domain: str) -> Optional[Dict[str, Union[str, None]]]:
        """提供字典形式的WHOIS内容

        Args:
            domain: 要查询的域名，如"x.com"

        Returns:
            包含WHOIS信息的字典，格式为{
                'expire_time': str,
            } 或 None（当无法获取有效信息时）

        Examples:
            >>> get_domain_whois("example.com")
            {
                'expire_time': '2031-01-08T09:59:51Z',
            }
        """
        whois_data = self.get_domain_whois(domain)
        expire_time = whois_data.get('expire_time', '')
        return epochExpire(expire_time)

    def update_whois_list_txt() -> str:
        """从GitHub获取最新的whois-servers.txt文件。

        该方法会从GitHub仓库获取whois服务器列表文件。提供了两个备用源：
        1. GitHub页面源
        2. GitHub raw文件源

        Returns:
            str: 获取到的whois-servers.txt文件。

        Note:
            备用源：
            - https://github.com/WooMai/whois-servers/blob/master/list.txt
            - https://raw.githubusercontent.com/WooMai/whois-servers/master/list.txt
            - https://www.nirsoft.net/whois-servers.txt

        Example:
            >>> update_whois_list_txt()
        """

        url = 'https://raw.gitmirror.com/WooMai/whois-servers/master/list.txt'
        response = requests.get(url, timeout=3)
        if response.ok:
            with io.open("/app/wingman/domain/whois-servers.txt", 'w', encoding='utf-8') as f:
                f.write(response.text)


if __name__ == '__main__':
    whois = CustomWhois()
    # print("Domain WHOIS Info:")
    # print(whois.get_domain_whois('x.com'))
    print(whois.get_domain_expire('yxbkj.cn'))
